import argparse
import random
import csv

def parse_args():
    parser = argparse.ArgumentParser(description="从tsv1中无重复随机采样指定数目的不在tsv2中的数据到tsv3")
    parser.add_argument("--tsv1", required=True, help="输入的tsv1文件路径")
    parser.add_argument("--tsv2", required=True, help="输入的tsv2文件路径")
    parser.add_argument("--tsv3", required=True, help="输出的tsv3文件路径")
    parser.add_argument("--sample_size", type=int, required=True, help="需要采样的记录数量")
    return parser.parse_args()

def read_tsv2_keys(tsv2_path):
    """
    读取tsv2文件中的键(chrom, pos, alignstrand, readname)，并存储到集合中
    """
    keys = set()
    with open(tsv2_path, "r") as file:
        reader = csv.reader(file, delimiter="\t")
        for line in reader:
            # 创建键：chrom, pos, alignstrand, readname
            key = (line[0], line[1], line[2], line[3])
            keys.add(key)
    return keys

def find_candidate_lines(tsv1_path, tsv2_keys):
    """
    找到符合条件的行号
    """
    candidate_lines = []
    with open(tsv1_path, "r") as infile:
        reader = csv.reader(infile, delimiter="\t")
        for line_number, line in enumerate(reader):
            # 创建键：chrom, pos, alignstrand, readname
            key = (line[0], line[1], line[2], line[4])
            # 如果该行不在tsv2的键集合中，则记录行号
            if key not in tsv2_keys:
                candidate_lines.append(line_number)
    return candidate_lines

def write_selected_lines(tsv1_path, selected_line_numbers, tsv3_path):
    """
    写出被选中的行
    """
    with open(tsv1_path, "r") as infile, open(tsv3_path, "w", newline="") as outfile:
        reader = csv.reader(infile, delimiter="\t")
        writer = csv.writer(outfile, delimiter="\t")
        
        for line_number, line in enumerate(reader):
            if line_number in selected_line_numbers:
                writer.writerow(line)

def main():
    args = parse_args()
    
    # 读取tsv2中的键
    tsv2_keys = read_tsv2_keys(args.tsv2)
    
    # 找到符合条件的行号
    candidate_lines = find_candidate_lines(args.tsv1, tsv2_keys)
    
    # 随机抽样行号
    if len(candidate_lines) > args.sample_size:
        selected_lines = random.sample(candidate_lines, args.sample_size)
    else:
        selected_lines = candidate_lines  # 如果候选不足，返回所有行号
    
    # 写出被选中的行
    write_selected_lines(args.tsv1, set(selected_lines), args.tsv3)

if __name__ == "__main__":
    main()
